/*

 * Copyright (c) 2017-2018, FastDev 刘强 (fastdev@163.com).

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package of.hotswapinject.agent;

import of.hotswapinject.core.common.ConfigUtil;
import of.hotswapinject.core.common.Const;
import of.hotswapinject.core.common.IPluginQueryable;
import of.hotswapinject.core.common.RuntimeContext;
import of.hotswapinject.core.common.Util;
import of.hotswapinject.injection.MethodHook;
import of.hotswapinject.injection.MethodHookTransformer;
import of.hotswapinject.core.plugin.IMethodFilter;
import of.hotswapinject.core.search.ISearch;
import com.google.common.reflect.ClassPath;
import org.apache.commons.cli.CommandLine;
import org.objectweb.asm.Opcodes;

import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.util.*;

public class TransformerInstaller {
     private boolean search(CommandLine cmdLine, String option, String name, List<String> matchedClasses, ISearch iSearch) {
        String value = cmdLine.getOptionValue(option);
        if(iSearch.filter(name, value)) {
            matchedClasses.add(name);
            return true;
        }
        return false;
    }

    private List<String> getAll() throws IOException {
        List<String> matchedClasses = new ArrayList<String>();
        ClassPath cp = ClassPath.from(Thread.currentThread().getContextClassLoader());
        for(ClassPath.ClassInfo ci : cp.getAllClasses()) {
            matchedClasses.add(ci.getName());
        }
        return matchedClasses;
    }

    private static void setPluginQueryable(CommandLine cmdLine) {
        if(RuntimeContext.getContext().getPluginQueryable() == null) {
            RuntimeContext.getContext().setPluginQueryable(new IPluginQueryable() {
                @Override
                public IMethodFilter getMethodFilter(String key) {
                    return MethodHook.getMethodFilter(key);
                }

                @Override
                public String getConfig(String config) {
                    return ConfigUtil.getConfig(config);
                }

                @Override
                public boolean isException(int opCode) {
                    return opCode == Opcodes.ATHROW;
                }
            });
        }
    }

    public void install(PluginInstallContext context, CommandLine cmdLine, Instrumentation inst) throws IOException {
        setPluginQueryable(cmdLine);
        List<Class> retransClassList = Arrays.asList(inst.getAllLoadedClasses());

        List<String> matchedClasses = new ArrayList<String>();
        boolean hasSearchPluginMatched = false;
        for(Map.Entry<String, ISearch> entry : context.searchMap.entrySet()) {
            if(cmdLine.hasOption(entry.getKey())) {
                List<Class> newList = new ArrayList<Class>();
                for(Class cc : retransClassList) {
                    if (search(cmdLine, entry.getKey(), cc.getName(), matchedClasses, entry.getValue())) {
                        newList.add(cc);
                    }
                }
                retransClassList = newList;
                if(!cmdLine.hasOption("filterLoaded"))
                {
                    matchedClasses = new ArrayList<String>();
                    for(String cName : getAll()) {
                        if (search(cmdLine, entry.getKey(), cName, matchedClasses, entry.getValue())) {
                        }
                    }
                }
                hasSearchPluginMatched = true;
                break;
            }
        }

        boolean hasListMatch = false;
        if(!hasSearchPluginMatched) {
            if(cmdLine.hasOption("filterLoaded")) {
                for(Class cc : retransClassList) {
                    matchedClasses.add(cc.getName());
                }
                hasListMatch = true;
            }
            else if(cmdLine.hasOption("list")) {
                matchedClasses = getAll();
                hasListMatch = true;
            }
        }

        if(!hasListMatch && !hasSearchPluginMatched) {
            Util.trace("no class selected!");
            return;
        }

        if(cmdLine.hasOption("output")) {
            String outputPath = cmdLine.getOptionValue("output");
            Util.saveFile(outputPath, String.join(Const.NEWLINE, matchedClasses));
            System.out.println("----------------------Search list to:" + outputPath + "----------------------");
        }
        else {
            System.out.println("----------------------Search list result:----------------------");
            System.out.println(String.join(Const.NEWLINE, matchedClasses));
        }

        List<IMethodFilter> methodFiltes = context.getMethodInjectors(cmdLine);
        if(methodFiltes == null || methodFiltes.size() == 0) {
            return;
        }

        if(methodFiltes.size() > 1) {
            throw new RuntimeException("can't apply more than one method of.hotswapinject.core.plugin!");
        }
        inst.addTransformer(new MethodHookTransformer(methodFiltes.get(0), matchedClasses, cmdLine.getOptionValue("dump")), true);
        try {
            Class[] modifiableClasses = getRetransClassList(inst, retransClassList);
            if(modifiableClasses == null || modifiableClasses.length == 0) {
                return;
            }
            inst.retransformClasses(modifiableClasses);
        } catch (UnmodifiableClassException e) {
            e.printStackTrace();
        }
    }

    private Class[] getRetransClassList(Instrumentation inst, List<Class> retransClassList) {
        if(retransClassList.size() <= 0) {
            return null;
        }
        List<Class> modifiableClasses = new ArrayList<Class>(retransClassList.size());
         for(Class clazz : retransClassList) {
             if (inst.isModifiableClass(clazz) && !clazz.getName().startsWith("java.lang.invoke.LambdaForm")) {
                 modifiableClasses.add(clazz);
             }
         }
         if(modifiableClasses.size() <= 0) {
             return null;
         }
         return modifiableClasses.toArray(new Class[modifiableClasses.size()]);
    }
}
